#!/usr/bin/python3
# coding:utf-8

import rospy
from cv_bridge import CvBridge
from control.msg import detect_res
from sensor_msgs.msg import Image, CameraInfo
import cv2
import numpy as np

class ColorDetectNode:
    def __init__(self):
        rospy.init_node('color_detect_2_node', anonymous = True)          #创建节点
        self.rate = rospy.Rate(20)
        self.bridge = CvBridge()

        rospy.Subscriber("/protobot_camera_2/camera_info",CameraInfo,self.camera_info_cb)
        rospy.Subscriber("/protobot_camera_2/image_raw",Image,self.detect_cb)
        
        self.color_dist = {'red': {'Lower': np.array([0, 60, 60]), 'Upper': np.array([6, 255, 255])},
                        'blue': {'Lower': np.array([115, 200, 200]), 'Upper': np.array([124, 255, 255])},
                        'green': {'Lower': np.array([35, 43, 35]), 'Upper': np.array([90, 255, 255])},
                        }
        #pub1 = rospy.Publisher(“/topic_name”, message_type, queue_size=size)
        #detect_res float32[] pos uint8 shape
        self.detect_res_pub = rospy.Publisher("/detect_res_2",detect_res,queue_size=1)
        self.img_res_pub = rospy.Publisher("/img_res_2",Image,queue_size=1)

        # wait to modify
        self.RT_mat = np.array(np.eye(4))
        rospy.loginfo("Color detect 2 is ready.")

    def camera_info_cb(self,data):
        self.K = np.array(data.K).reshape((3,3))
        self.P = np.matrix(np.array(data.P).reshape((3,4))).I
        self.f_x = self.K[0,0]
        self.f_y = self.K[1,1]

    def detect_cb(self,data):
        frame = self.bridge.imgmsg_to_cv2(data)
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        inRange_hsv = cv2.inRange(hsv, self.color_dist['red']['Lower'], self.color_dist['red']['Upper'])
        cnts = cv2.findContours(inRange_hsv.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
        if cnts:
            #rospy.loginfo("detect red")
            c = max(cnts, key=cv2.contourArea)
            shape = 0
            if len(cv2.approxPolyDP(c,0.01*cv2.arcLength(c,True),True))<=4:
                shape = 1
                #rospy.loginfo("detect red square")
            else:
                shape = 2
                #rospy.loginfo("detect red circle")
            rect = cv2.minAreaRect(c)
            box = cv2.boxPoints(rect)
            target = np.sum(box,axis=0)/4
            target = np.array([[target[0]],
                            [target[1]],
                            [1]
                            ])
            temp_vec = np.dot(self.P,target)
            temp_vec = np.array([[temp_vec[0]],
                                [temp_vec[1]],
                                [temp_vec[2]],
                                [0.0]])
            target_global = np.array(np.dot(self.RT_mat,temp_vec))
            temp_detect_res = detect_res()
            for i in target_global:
                temp_detect_res.pos.append(i)
            temp_detect_res.shape = shape
            self.detect_res_pub.publish(temp_detect_res)
            frame_rgb = cv2.circle(frame_rgb, (int(target[0]),int(target[1])), 1, (0, 0, 255), -1)
            frame_rgb = cv2.drawContours(frame_rgb, [np.int0(box)], -1, (0, 0, 255), 1)

        inRange_hsv = cv2.inRange(hsv, self.color_dist['blue']['Lower'], self.color_dist['blue']['Upper'])
        cnts = cv2.findContours(inRange_hsv.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
        if cnts:
            #rospy.loginfo("detect blue")
            c = max(cnts, key=cv2.contourArea)
            shape = 0
            if len(cv2.approxPolyDP(c,0.01*cv2.arcLength(c,True),True))<=4:
                shape = 1
                #rospy.loginfo("detect blue square")
            else:
                shape = 2
                #rospy.loginfo("detect blue circle")
            rect = cv2.minAreaRect(c)
            box = cv2.boxPoints(rect)
            target = np.sum(box,axis=0)/4
            target = np.array([[target[0]],
                            [target[1]],
                            [1.0]])
            temp_vec = np.dot(self.P,target)
            temp_vec = np.array([[temp_vec[0]],
                                [temp_vec[1]],
                                [temp_vec[2]],
                                [0.0]])
            target_global = np.array(np.dot(self.RT_mat,temp_vec))
            temp_detect_res = detect_res()
            for i in target_global:
                temp_detect_res.pos.append(i)
            temp_detect_res.shape = shape
            self.detect_res_pub.publish(temp_detect_res)
            frame_rgb = cv2.circle(frame_rgb, (int(target[0]),int(target[1])), 1, (255, 0, 0), -1)
            frame_rgb = cv2.drawContours(frame_rgb, [np.int0(box)], -1, (255, 0, 0), 1)

        self.img_res_pub.publish(self.bridge.cv2_to_imgmsg(frame_rgb,'rgb8'))


    def MainLoop(self):
        while not rospy.is_shutdown():
            self.rate.sleep()
            
if __name__ == '__main__':
    colordetectNode = ColorDetectNode()
    colordetectNode.MainLoop()
